# https://blog.csdn.net/bananapai/article/details/145736300

import torch
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import Compose, Resize, ToTensor

from Alexnet2MNIST_model import AlexNet

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
ds = MNIST(root='./data', train=False, download=False,
           transform=Compose([Resize((224, 224), antialias=True), ToTensor()]),
           target_transform=lambda label: torch.tensor(label))
dl = DataLoader(ds, batch_size=1000, shuffle=False)
model = AlexNet()
model.to(device)
model.load_state_dict(torch.load('weights/model_detect_numbers.pt', weights_only=True))
model.eval()

correct_count = 0
for i, (inputs, labels) in enumerate(dl):
    with torch.inference_mode():
        y = model(inputs)
    y = torch.nn.functional.softmax(y, dim=-1)
    # 求最大值索引
    idx = y.argmax(-1)
    correct_count += (idx == labels).short().sum().item()
    print('准确的数量', correct_count)

print(f'准确率: {correct_count / len(ds) * 100:.2f}%')
print("Alexnet2MNIST_test.py", 444)